from flask import Flask, render_template, request, jsonify, redirect, url_for, flash, session
from flask_sqlalchemy import SQLAlchemy
from flask_login import LoginManager, UserMixin, login_user, logout_user, login_required, current_user
from flask_wtf import FlaskForm
from wtforms import StringField, PasswordField, TextAreaField, SelectField, SubmitField
from wtforms.validators import DataRequired, Length, Email
from flask_cors import CORS
import os
import psutil
import requests
import json
from datetime import datetime
import threading
import time
import sys

# 设置输出编码
if sys.platform.startswith('win'):
    import codecs
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.detach())
    sys.stderr = codecs.getwriter('utf-8')(sys.stderr.detach())

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key-here'
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///network_system.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

db = SQLAlchemy(app)
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'login'
CORS(app)

# 数据库模型
class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    email = db.Column(db.String(120), unique=True, nullable=False)
    password = db.Column(db.String(120), nullable=False)
    role = db.Column(db.String(20), default='recon')  # 用户的默认/授权角色（现用作授权白名单）
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

# 侦察数据表
class ReconData(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(200), nullable=False)
    target_system = db.Column(db.String(100), nullable=False)
    data_type = db.Column(db.String(50), nullable=False)  # network, system, vulnerability, etc.
    content = db.Column(db.Text, nullable=False)
    priority = db.Column(db.String(20), default='medium')  # low, medium, high, critical
    status = db.Column(db.String(20), default='pending')  # pending, analyzed, processed
    created_by = db.Column(db.Integer, db.ForeignKey('user.id'))
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

# 分析结果表
class AnalysisResult(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    recon_data_id = db.Column(db.Integer, db.ForeignKey('recon_data.id'))
    analysis_type = db.Column(db.String(50), nullable=False)  # threat, vulnerability, opportunity
    findings = db.Column(db.Text, nullable=False)
    recommendations = db.Column(db.Text, nullable=False)
    priority_targets = db.Column(db.Text)  # JSON格式存储目标列表
    suggested_actions = db.Column(db.Text)  # JSON格式存储建议行动
    confidence_level = db.Column(db.String(20), default='medium')  # low, medium, high
    status = db.Column(db.String(20), default='pending')  # pending, reviewed, approved
    created_by = db.Column(db.Integer, db.ForeignKey('user.id'))
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

# 决策指令表
class DecisionOrder(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    analysis_id = db.Column(db.Integer, db.ForeignKey('analysis_result.id'))
    order_type = db.Column(db.String(50), nullable=False)  # reconnaissance, analysis, execution
    target = db.Column(db.String(200), nullable=False)
    objective = db.Column(db.Text, nullable=False)
    instructions = db.Column(db.Text, nullable=False)
    priority = db.Column(db.String(20), default='medium')
    deadline = db.Column(db.DateTime)
    status = db.Column(db.String(20), default='pending')  # pending, executing, completed, failed
    assigned_to = db.Column(db.Integer, db.ForeignKey('user.id'))
    created_by = db.Column(db.Integer, db.ForeignKey('user.id'))
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

# 执行结果表
class ExecutionResult(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    order_id = db.Column(db.Integer, db.ForeignKey('decision_order.id'))
    result_type = db.Column(db.String(50), nullable=False)  # success, partial, failed
    description = db.Column(db.Text, nullable=False)
    evidence = db.Column(db.Text)  # 执行证据或截图
    impact_assessment = db.Column(db.Text)
    lessons_learned = db.Column(db.Text)
    status = db.Column(db.String(20), default='pending')  # pending, reviewed
    created_by = db.Column(db.Integer, db.ForeignKey('user.id'))
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

# 系统日志表
class SystemLog(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    timestamp = db.Column(db.DateTime, default=datetime.utcnow)
    level = db.Column(db.String(20), nullable=False)
    message = db.Column(db.Text, nullable=False)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'))
    action_type = db.Column(db.String(50))  # recon, analysis, decision, execution

# 表单类
class LoginForm(FlaskForm):
    username = StringField('用户名', validators=[DataRequired()])
    password = PasswordField('密码', validators=[DataRequired()])
    submit = SubmitField('登录')

class ReconDataForm(FlaskForm):
    title = StringField('侦察标题', validators=[DataRequired(), Length(min=1, max=200)])
    target_system = StringField('目标系统', validators=[DataRequired(), Length(min=1, max=100)])
    data_type = SelectField('数据类型', choices=[
        ('network', '网络信息'),
        ('system', '系统信息'),
        ('vulnerability', '漏洞信息'),
        ('user', '用户信息'),
        ('service', '服务信息'),
        ('other', '其他')
    ])
    content = TextAreaField('侦察内容', validators=[DataRequired()])
    priority = SelectField('优先级', choices=[
        ('low', '低'),
        ('medium', '中'),
        ('high', '高'),
        ('critical', '紧急')
    ])
    submit = SubmitField('提交侦察数据')

class AnalysisForm(FlaskForm):
    analysis_type = SelectField('分析类型', choices=[
        ('threat', '威胁分析'),
        ('vulnerability', '漏洞分析'),
        ('opportunity', '机会分析'),
        ('comprehensive', '综合分析')
    ])
    findings = TextAreaField('分析发现', validators=[DataRequired()])
    recommendations = TextAreaField('建议措施', validators=[DataRequired()])
    priority_targets = TextAreaField('优先目标')
    suggested_actions = TextAreaField('建议行动')
    confidence_level = SelectField('置信度', choices=[
        ('low', '低'),
        ('medium', '中'),
        ('high', '高')
    ])
    submit = SubmitField('提交分析结果')

class DecisionForm(FlaskForm):
    order_type = SelectField('指令类型', choices=[
        ('reconnaissance', '侦察指令'),
        ('analysis', '分析指令'),
        ('execution', '执行指令')
    ])
    target = StringField('目标', validators=[DataRequired(), Length(min=1, max=200)])
    objective = TextAreaField('目标描述', validators=[DataRequired()])
    instructions = TextAreaField('具体指令', validators=[DataRequired()])
    priority = SelectField('优先级', choices=[
        ('low', '低'),
        ('medium', '中'),
        ('high', '高'),
        ('critical', '紧急')
    ])
    assigned_to = SelectField('分配给', coerce=int)
    submit = SubmitField('下达指令')

class ExecutionForm(FlaskForm):
    result_type = SelectField('执行结果', choices=[
        ('success', '成功'),
        ('partial', '部分成功'),
        ('failed', '失败')
    ])
    description = TextAreaField('执行描述', validators=[DataRequired()])
    evidence = TextAreaField('执行证据')
    impact_assessment = TextAreaField('影响评估')
    lessons_learned = TextAreaField('经验教训')
    submit = SubmitField('提交执行结果')

@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))

# 辅助函数
def log_system_event(level, message, action_type=None):
    log = SystemLog(
        level=level, 
        message=message, 
        user_id=current_user.id if current_user.is_authenticated else None,
        action_type=action_type
    )
    db.session.add(log)
    db.session.commit()

def get_active_role():
    return session.get('active_role')

def set_active_role(role_name):
    session['active_role'] = role_name

def clear_active_role():
    session.pop('active_role', None)

def role_required(role_name):
    def decorator(func):
        from functools import wraps
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not current_user.is_authenticated:
                return redirect(url_for('login'))
            active = get_active_role()
            if active != role_name:
                flash('请先选择并通过该身份的权限验证')
                return redirect(url_for('select_role'))
            return func(*args, **kwargs)
        return wrapper
    return decorator

def get_role_dashboard():
    """根据当前已启用的身份返回对应仪表板"""
    if not current_user.is_authenticated:
        return redirect(url_for('login'))
    
    role = get_active_role()
    if not role:
        return redirect(url_for('select_role'))
    if role == 'recon':
        return redirect(url_for('recon_dashboard'))
    elif role == 'analyst':
        return redirect(url_for('analyst_dashboard'))
    elif role == 'decision':
        return redirect(url_for('decision_dashboard'))
    elif role == 'executor':
        return redirect(url_for('executor_dashboard'))
    else:
        return redirect(url_for('login'))

# 路由
@app.route('/')
def index():
    if current_user.is_authenticated:
        return get_role_dashboard()
    return render_template('index.html')

@app.route('/login', methods=['GET', 'POST'])
def login():
    if current_user.is_authenticated:
        return get_role_dashboard()
    
    form = LoginForm()
    if form.validate_on_submit():
        user = User.query.filter_by(username=form.username.data).first()
        if user and user.password == form.password.data:  # 简化密码验证
            login_user(user)
            clear_active_role()
            log_system_event('info', f'用户 {user.username} 登录系统', 'login')
            return redirect(url_for('select_role'))
        flash('用户名或密码错误')
    
    return render_template('login.html', form=form)

@app.route('/logout')
@login_required
def logout():
    log_system_event('info', f'用户 {current_user.username} 退出系统', 'logout')
    logout_user()
    clear_active_role()
    return redirect(url_for('index'))

@app.context_processor
def inject_active_role():
    return {'active_role': get_active_role()}

@app.route('/select_role', methods=['GET', 'POST'])
@login_required
def select_role():
    if request.method == 'POST':
        selected_role = request.form.get('role')
        role_password = request.form.get('role_password')
        if selected_role not in ['recon', 'analyst', 'decision', 'executor']:
            flash('身份选择无效')
            return redirect(url_for('select_role'))
        # 二次权限验证（再次输入账号密码）
        if not role_password or current_user.password != role_password:
            flash('身份验证失败，请输入正确的密码以启用该身份')
            return redirect(url_for('select_role'))
        # 放宽：所有用户均可启用四种身份
        set_active_role(selected_role)
        log_system_event('info', f"用户 {current_user.username} 启用身份: {selected_role}", selected_role)
        return get_role_dashboard()
    return render_template('select_role.html')

# 侦察员页面
@app.route('/recon')
@login_required
@role_required('recon')
def recon_dashboard():
    
    # 获取用户提交的侦察数据
    recon_data = ReconData.query.filter_by(created_by=current_user.id).order_by(ReconData.created_at.desc()).all()
    
    return render_template('recon_dashboard.html', recon_data=recon_data)

@app.route('/recon/add', methods=['GET', 'POST'])
@login_required
@role_required('recon')
def add_recon_data():
    
    form = ReconDataForm()
    if form.validate_on_submit():
        recon_data = ReconData(
            title=form.title.data,
            target_system=form.target_system.data,
            data_type=form.data_type.data,
            content=form.content.data,
            priority=form.priority.data,
            created_by=current_user.id
        )
        db.session.add(recon_data)
        db.session.commit()
        
        log_system_event('info', f'侦察员 {current_user.username} 提交了新的侦察数据: {form.title.data}', 'recon')
        flash('侦察数据提交成功！')
        return redirect(url_for('recon_dashboard'))
    
    return render_template('add_recon_data.html', form=form)

# 分析员页面
@app.route('/analyst')
@login_required
@role_required('analyst')
def analyst_dashboard():
    
    # 获取待分析的侦察数据
    pending_recon = ReconData.query.filter_by(status='pending').order_by(ReconData.created_at.desc()).all()
    # 获取用户的分析结果
    analysis_results = AnalysisResult.query.filter_by(created_by=current_user.id).order_by(AnalysisResult.created_at.desc()).all()
    
    return render_template('analyst_dashboard.html', pending_recon=pending_recon, analysis_results=analysis_results)

@app.route('/analyst/analyze/<int:recon_id>', methods=['GET', 'POST'])
@login_required
@role_required('analyst')
def analyze_data(recon_id):
    
    recon_data = ReconData.query.get_or_404(recon_id)
    form = AnalysisForm()
    
    if form.validate_on_submit():
        analysis = AnalysisResult(
            recon_data_id=recon_id,
            analysis_type=form.analysis_type.data,
            findings=form.findings.data,
            recommendations=form.recommendations.data,
            priority_targets=form.priority_targets.data,
            suggested_actions=form.suggested_actions.data,
            confidence_level=form.confidence_level.data,
            created_by=current_user.id
        )
        db.session.add(analysis)
        
        # 更新侦察数据状态
        recon_data.status = 'analyzed'
        db.session.commit()
        
        log_system_event('info', f'分析员 {current_user.username} 完成了对侦察数据 "{recon_data.title}" 的分析', 'analysis')
        flash('分析结果提交成功！')
        return redirect(url_for('analyst_dashboard'))
    
    return render_template('analyze_data.html', form=form, recon_data=recon_data)

# 决策员页面
@app.route('/decision')
@login_required
@role_required('decision')
def decision_dashboard():
    
    # 获取待审核的分析结果
    pending_analysis = AnalysisResult.query.filter_by(status='pending').order_by(AnalysisResult.created_at.desc()).all()
    # 获取用户下达的指令
    orders = DecisionOrder.query.filter_by(created_by=current_user.id).order_by(DecisionOrder.created_at.desc()).all()
    # 获取执行结果
    execution_results = ExecutionResult.query.join(DecisionOrder).filter(DecisionOrder.created_by==current_user.id).order_by(ExecutionResult.created_at.desc()).all()
    
    return render_template('decision_dashboard.html', 
                         pending_analysis=pending_analysis, 
                         orders=orders, 
                         execution_results=execution_results)

@app.route('/decision/create_order', methods=['GET', 'POST'])
@login_required
@role_required('decision')
def create_order():
    
    form = DecisionForm()
    # 放寬：所有用戶均可被指派（如需依任務類型細分，之後再調整）
    form.assigned_to.choices = [(u.id, f"{u.username}") for u in User.query.all()]
    
    if form.validate_on_submit():
        order = DecisionOrder(
            order_type=form.order_type.data,
            target=form.target.data,
            objective=form.objective.data,
            instructions=form.instructions.data,
            priority=form.priority.data,
            assigned_to=form.assigned_to.data,
            created_by=current_user.id
        )
        db.session.add(order)
        db.session.commit()
        
        assigned_user = User.query.get(form.assigned_to.data)
        log_system_event('info', f'决策员 {current_user.username} 向 {assigned_user.username} 下达了指令: {form.target.data}', 'decision')
        flash('指令下达成功！')
        return redirect(url_for('decision_dashboard'))
    
    return render_template('create_order.html', form=form)

# 执行员页面
@app.route('/executor')
@login_required
@role_required('executor')
def executor_dashboard():
    
    # 获取分配给当前用户的指令
    assigned_orders = DecisionOrder.query.filter_by(assigned_to=current_user.id, status='pending').order_by(DecisionOrder.created_at.desc()).all()
    # 获取用户提交的执行结果
    execution_results = ExecutionResult.query.filter_by(created_by=current_user.id).order_by(ExecutionResult.created_at.desc()).all()
    
    return render_template('executor_dashboard.html', assigned_orders=assigned_orders, execution_results=execution_results)

@app.route('/executor/execute/<int:order_id>', methods=['GET', 'POST'])
@login_required
@role_required('executor')
def execute_order(order_id):
    
    order = DecisionOrder.query.get_or_404(order_id)
    if order.assigned_to != current_user.id:
        flash('权限不足')
        return redirect(url_for('executor_dashboard'))
    
    form = ExecutionForm()
    
    if form.validate_on_submit():
        execution = ExecutionResult(
            order_id=order_id,
            result_type=form.result_type.data,
            description=form.description.data,
            evidence=form.evidence.data,
            impact_assessment=form.impact_assessment.data,
            lessons_learned=form.lessons_learned.data,
            created_by=current_user.id
        )
        db.session.add(execution)
        
        # 更新指令状态
        order.status = 'completed'
        db.session.commit()
        
        log_system_event('info', f'执行员 {current_user.username} 完成了指令: {order.target}', 'execution')
        flash('执行结果提交成功！')
        return redirect(url_for('executor_dashboard'))
    
    return render_template('execute_order.html', form=form, order=order)

# 系统日志页面
@app.route('/logs')
@login_required
def logs():
    page = request.args.get('page', 1, type=int)
    logs = SystemLog.query.order_by(SystemLog.timestamp.desc()).paginate(
        page=page, per_page=20, error_out=False)
    return render_template('logs.html', logs=logs)

if __name__ == '__main__':
    with app.app_context():
        db.create_all()
        
        # 创建默认用户
        default_users = [
            {'username': 'recon1', 'email': 'recon1@example.com', 'password': 'recon123', 'role': 'recon'},
            {'username': 'analyst1', 'email': 'analyst1@example.com', 'password': 'analyst123', 'role': 'analyst'},
            {'username': 'decision1', 'email': 'decision1@example.com', 'password': 'decision123', 'role': 'decision'},
            {'username': 'executor1', 'email': 'executor1@example.com', 'password': 'executor123', 'role': 'executor'},
        ]
        
        for user_data in default_users:
            if not User.query.filter_by(username=user_data['username']).first():
                user = User(**user_data)
                db.session.add(user)
        
        db.session.commit()
        print("默认用户账户已创建:")
        print("侦察员: recon1/recon123")
        print("分析员: analyst1/analyst123")
        print("决策员: decision1/decision123")
        print("执行员: executor1/executor123")
    
    app.run(debug=True, host='0.0.0.0', port=5000)
